/**
 * @file algo_Filter.c
 * @brief Software implementation of EKF for 51 microcontroller.
 * @author David Xu
 * @version 1.0.1
 * @date 2023-12-06
 * @note none
 */

#include "algo_ekf.h"
#include "algo_matrix.h"

/**
 * @brief State-space model function.
 * @param[in,out] kf Pointer to the EKF filter structure.
 * @param[in] u Control input.
 * @param[in] dt Time step.
 * @note kf Pointer use to access the state vector X[in],
 * state derivative Xdot[out], and state transition matrix F[out].
 */
void ekf_sfun(struct ekf_filter *kf, const float *u, const float dt) {
    // Access the state vector X, state derivative Xdot, and state transition matrix F
    // directly from the ekf_filter structure pointed by kf.

    // Assuming a simple linear system with state transition matrix A and control input matrix B
    // Xdot = A * X + B * u
    // F is the Jacobian of Xdot with respect to X, which is A in a linear system

    // Define the state transition matrix A (example for a 2D system)
    Matrix_fp16_t A;
    Matrix_fp16_t B;
    Matrix_fp16_t Bu;
    qFP16_mat_setup(&A, 0);    // Initialize A with zeros
    A.Mdat[0][0] = qFP16_Constant(1.0f);
    A.Mdat[0][1] = qFP16_FloatToFP(dt);
    A.Mdat[1][0] = qFP16_Constant(0.0f);
    A.Mdat[1][1] = qFP16_Constant(1.0f);

    // Define the control input matrix B (example for a 2D system)
    qFP16_mat_setup(&B, 0);    // Initialize B with zeros
    B.Mdat[1][0] = qFP16_FloatToFP(dt);

    // Compute Xdot
    qFP16_mat_mult(&kf->Xdot, &A, &kf->X);       // Xdot = A * X
    qFP16_mat_scal_mult(&Bu, &B, *u);            // Bu = B * u
    qFP16_mat_add(&kf->Xdot, &kf->Xdot, &Bu);    // Xdot = Xdot + Bu

    // Copy A to F since it's the Jacobian of Xdot with respect to X in a linear system
    qFP16_mat_copy(&kf->F, &A);

    // TODO: Define the state transition model here.
    // For example, for a simple linear system:
    // Xdot[0] = F[0] * X[0] + F[1] * X[1] + u[0];
    // Xdot[1] = F[2] * X[0] + F[3] * X[1] + u[1];

    // TODO: Fill the state transition matrix F based on the system dynamics.
    // For a linear system, this would be constant. For a nonlinear system,
    // it would be the Jacobian of the state transition function with respect to the state
    // vector.
}

/**
 * @brief Measurement model function.
 * @param[in,out] kf Pointer to the EKF filter structure.
 * @param[in] y Measurement vector.
 * @note kf Pointer use to access the measurement matrix X[in], H[out] and measurement error err[out].
 */

void ekf_mfun(struct ekf_filter *kf, const Matrix_fp16_t *y) {
    // Access the state vector X, measurement matrix H, and measurement error directly
    // from the ekf_filter structure pointed by kf.

    // Assuming a simple linear measurement model with measurement matrix C
    // y = C * X
    // H is the Jacobian of y with respect to X, which is C in a linear system

    // Define the measurement matrix C (example for a 2D system)
    Matrix_fp16_t C;
    Matrix_fp16_t y_est;
    qFP16_mat_setup(&C, 0);    // Initialize C with zeros
    C.Mdat[0][0] = qFP16_Constant(1.0f);
    C.Mdat[1][1] = qFP16_Constant(1.0f);

    // Compute the measurement error
    qFP16_mat_mult(&y_est, &C, &kf->X);    // y_est = C * X
    qFP16_mat_sub(&kf->err, y, &y_est);    // err = y - y_est

    // Copy C to H since it's the Jacobian of y with respect to X in a linear system
    qFP16_mat_copy(&kf->H, &C);
    // TODO: Define the measurement model here.
    // For example, for a simple linear system:
    // err[0] = y[0] - (H[0] * X[0] + H[1] * X[1]);
    // err[1] = y[1] - (H[2] * X[0] + H[3] * X[1]);

    // TODO: Fill the measurement matrix H based on the system measurement model.
    // For a linear system, this would be constant. For a nonlinear system,
    // it would be the Jacobian of the measurement function with respect to the state
    // vector.
}


/*============================================================================*/

void ekf_init(struct ekf_filter *kf,
              const float       *initial_state,
              const float       *initial_P,
              const float       *Q,
              const float       *R) {
    // Set the state dimension and measurement dimension
    kf->state_dim   = STATE_DIM;
    kf->measure_dim = MEASURE_DIM;

    // Initialize the state vector X
    qFP16_array_to_matrix(&kf->X, initial_state);

    // Initialize the error covariance matrix P
    qFP16_array_to_matrix(&kf->P, initial_P);

    // Initialize the process noise covariance matrix Q
    qFP16_array_to_matrix(&kf->Q, Q);

    // Initialize the measurement noise covariance matrix R
    qFP16_array_to_matrix(&kf->R, R);

    // Initialize the state transition matrix F and measurement matrix H as zero matrices,
    // they will be updated in ekf_sfun and ekf_mfun
    qFP16_mat_zero(&kf->F);
    qFP16_mat_zero(&kf->H);

    // Initialize the Kalman gain matrix K as a zero matrix
    qFP16_mat_zero(&kf->K);

    // Initialize temporary matrices
    qFP16_mat_zero(&kf->Xdot);
    qFP16_mat_zero(&kf->Pdot);
    qFP16_mat_zero(&kf->E);
    qFP16_mat_zero(&kf->err); /* Initialize the measurement error matrix */
}
/*============================================================================*/
void ekf_filter_predict(struct ekf_filter *kf, float *u) {
    Matrix_fp16_t tmp1, tmp2, tmp3;
    float         dt = 0.01;    // Hardcoded time step

    // if ekf_sfun is fill kf->Xdot and kf->F
    ekf_sfun(kf, u, dt);

    // X = X + Xdot * dt
    qFP16_mat_scal_mult(&kf->Xdot, &kf->Xdot, dt);    // Xdot = Xdot * dt
    qFP16_mat_add(&kf->X, &kf->X, &kf->Xdot);         // X = X + Xdot

#ifdef EKF_UPDATE_CONTINUOUS
    /*
        continuous update
        Pdot = F * P + P * F' + Q
    */
    qFP16_mat_mult(&tmp1, &kf->F, &kf->P);
    qFP16_mat_transpose(&tmp2, &kf->F);
    qFP16_mat_mult(&tmp3, &kf->P, &tmp2);
    qFP16_mat_add(&kf->Pdot, &tmp1, &tmp3);
    qFP16_mat_add(&kf->Pdot, &kf->Pdot, &kf->Q);
#endif
#ifdef EKF_UPDATE_DISCRETE
    // Pdot = F * P * F' + Q
    qFP16_mat_mult(&tmp1, &kf->F, &kf->P);
    qFP16_mat_transpose(&tmp2, &kf->F);
    qFP16_mat_mult(&tmp3, &tmp1, &tmp2);
    qFP16_mat_add(&kf->Pdot, &tmp3, &kf->Q);
#endif

    // P = P + Pdot * dt
    qFP16_mat_scal_mult(&kf->Pdot, &kf->Pdot, dt);    // Pdot = Pdot * dt
    qFP16_mat_add(&kf->P, &kf->P, &kf->Pdot);         // P = P + Pdot
}


/*============================================================================*/

void ekf_filter_update(struct ekf_filter *kf, Matrix_fp16_t *y) {
    Matrix_fp16_t tmp1, tmp2, tmp3;
    Matrix_fp16_t I, Kt;

    /* update

       E = H * P * H' + R
       K = P * H' * inv(E)
       P = P - K * H * P
       X = X + K * err

    */

    // if ekf_mfun is fill err and kf->H
    ekf_mfun(kf, y);

    // E = H * P * H' + R
    qFP16_mat_mult(&tmp1, &kf->H, &kf->P);    // tmp1 = H * P
    qFP16_mat_transpose(&tmp2, &kf->H);       // tmp2 = H'
    qFP16_mat_mult(&tmp3, &tmp1, &tmp2);      // tmp3 = tmp1 * tmp2
    qFP16_mat_add(&kf->E, &tmp3, &kf->R);     // E = tmp3 + R

    // K = P * H' * inv(E)
    qFP16_mat_inverse(&tmp1, &kf->E);      // tmp1 = inv(E)
    qFP16_mat_mult(&Kt, &kf->P, &tmp2);    // Kt = P * H'
    qFP16_mat_mult(&kf->K, &Kt, &tmp1);    // K = Kt * inv(E)

    // P = (I - K * H) * P
    qFP16_mat_identity(&I);                   // I = identity matrix
    qFP16_mat_mult(&tmp1, &kf->K, &kf->H);    // tmp1 = K * H
    qFP16_mat_sub(&tmp2, &I, &tmp1);          // tmp2 = I - K * H
    qFP16_mat_mult(&kf->P, &tmp2, &kf->P);    // P = (I - K * H) * P

    // X = X + K * err
    qFP16_mat_mult(&tmp1, &kf->K, &kf->err);    // tmp1 = K * err
    qFP16_mat_add(&kf->X, &kf->X, &tmp1);       // X = X + K * err
}


#if 1

/* Define the initial state, error covariance, process noise, and measurement noise */
static const float initial_state[STATE_DIM]         = { 0.0f, 0.0f };
static const float initial_P[STATE_DIM * STATE_DIM] = { 1.0f, 0.0f, 0.0f, 1.0f };
static const float Q[STATE_DIM * STATE_DIM]         = { 0.1f, 0.0f, 0.0f, 0.1f };
static const float R[MEASURE_DIM * MEASURE_DIM]     = { 0.01f, 0.0f, 0.0f, 0.01f };

#define SIGNAL_SIZE 100 /* 定义信号数组的大小 */

/* 一个简单的函数，用于将rand()生成的随机数缩放到[50,150]的范围内 */
float generate_random_signal(void) {
    return 50.0f + (float)(rand() % 1001); /* 101是因为150-50+1=101，确保范围是[50,150] */
}
struct ekf_filter kf;
Matrix_fp16_t     y; /* 测量值 */
void              ekf_test_dynamic_filtering(void) {

    float u = 1.0f; /* 在这个简单的例子中，我们假设控制输入为1 */

    float dt = 0.01f; /* 时间步长 */
    int   i;

    /* 初始化随机数生成器 */
    srand(1234); /* 使用固定的种子 */

    /* 初始化EKF */
    ekf_init(&kf, initial_state, initial_P, Q, R);

    /* 生成随机信号并进行EKF滤波 */
    printf("Original Signal : Filtered Signal\n");
    for (i = 0; i < SIGNAL_SIZE; ++i) {
        /* 生成随机信号 */
        float original_signal = generate_random_signal();
        y.Mdat[0][0]          = original_signal; /* 假设我们的测量值就是信号值 */

        /* 执行EKF预测和更新步骤 */
        ekf_filter_predict(&kf, &u);
        ekf_filter_update(&kf, &y);

        /* 打印原始信号和滤波后的信号 */
        printf("%d : %ld\n", (int)original_signal, kf.X.Mdat[0][0]);
    }
}


#endif